/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: TbeGatherV2的实现文件
 * Author: zhangjian
 * Create: 2023-5-17
 */

#include "tbe_gatherv2_aicore.h"
#include <fstream>
#include <dlfcn.h>
#include <unistd.h>
#include <nlohmann/json.hpp>
#include <cstring>
#include <hccl/hccl_types.h>
#include <string>
#include <sstream>
#include "tbe_vector_reduce.h"
#include "mem_device_pub.h"
#include "runtime/rt.h"
#include "op_json_info.h"
#include "op_tiling.h"
#include "vector_tiling.h"
#include "adapter_rts.h"
#include "mmpa_api.h"
#include "gatherv2_tiling.h"
#include "gatherv2_tiling_rt2.h"
#ifndef HCCD
#include "platform/platform_info.h"
#include "exe_graph/runtime/tiling_context.h"
#include "exe_graph/runtime/tiling_parse_context.h"
#include "exe_graph/lowering/tiling_context_builder.h"
#include "exe_graph/lowering/tiling_parse_context_builder.h"
#include "graph/utils/op_desc_utils.h"
#include "graph/operator_factory.h"
#include "graph/operator_reg.h"
#include "graph/operator.h"
#include "vector_tiling_rt2.h"
#include "graph/ge_tensor.h"
#include "graph/op_desc.h"
#include "graph/utils/graph_utils.h"
#endif

using namespace std;

template <typename T, typename std::enable_if<(!std::is_array<T>::value), int32_t>::type = 0>
static void *CreateCompileInfo()
{
    auto p = new T();
    return p;
}
template <typename T> static void DeleteCompileInfo(void * const obj)
{
    return;
}

namespace TbeReduce {

constexpr u32 MEMSET_GATHER_TILING_DATA_MAX_SZIE = 1000;
constexpr u32 MEMSET_OVERFLOW_SZIE = 32;
constexpr u32 UB_BLOCK_SIZE = 32;
constexpr u32 kMAX_TILING_DATA_SIZE = 16UL * 1024UL;
constexpr u32 kMAX_WORKSPACE_COUNT = 16;
constexpr size_t INPUT_SIZE = 2;
fe::PlatFormInfos g_gatherPlatFormInfos;

TbeGatherV2::TbeGatherV2()
{
}

TbeGatherV2::~TbeGatherV2()
{
    if (overFlowPtr_ != nullptr) {
        hrtFree(overFlowPtr_);
        overFlowPtr_ = nullptr;
    }
}

HcclResult TbeGatherV2::Init() // 初始化
{
    CHK_PRT_RET(isInit_, HCCL_DEBUG("TbeGatherV2 had been initialized"), HCCL_SUCCESS);
    CHK_RET(hrtGetDeviceType(deviceType_));
    switch (deviceType_) {
        case DevType::DEV_TYPE_910:
            InitOpInfoMap910A();
            break;
        case DevType::DEV_TYPE_310P3:
            InitOpInfoMap310P3();
            break;
        case DevType::DEV_TYPE_910B:
        case DevType::DEV_TYPE_910_93:
            InitOpInfoMap910B();
            GetCoreNum();
            break;
        default:
            HCCL_ERROR("[TbeGatherV2][Init]devType[%u] is not supported", deviceType_);
            return HCCL_E_NOT_SUPPORT;
    }
    tilingDataHostPtr_ = new(nothrow) char[MEMSET_GATHER_TILING_DATA_MAX_SZIE]; // tiling内存申请
    CHK_PTR_NULL(tilingDataHostPtr_);
    CHK_RET(LoadOpBinary()); // 加载二进制文件
    isInit_ = true;
    return HCCL_SUCCESS;
}

#ifndef HCCD
// rt1.0
HcclResult TbeGatherV2::GetTilingRunInfo(nlohmann::json &opDescInfo, nlohmann::json &opTilingInfo,
    u64 inputCount, u64 outputCount, u32 valueDim, HcclDataType dataType, OpRunInfo &runInfo)
{
    TeOpTensor inputTensor;
    inputTensor.dtype = "float32";
    inputTensor.shape = { static_cast<s64>(inputCount), valueDim };
    inputTensor.oriShape = { static_cast<s64>(inputCount), valueDim };
    inputTensor.format = "ND";
    inputTensor.oriFormat = "ND";

    TeOpTensorArg inputTilingArg;
    inputTilingArg.tensor.push_back(inputTensor);
    inputTilingArg.argType = TensorArgType::TA_SINGLE;

    TeOpTensor indicesInputTensor;

    if (dataType == HCCL_DATA_TYPE_INT32) {
        indicesInputTensor.dtype = "int32";
    } else if (dataType == HCCL_DATA_TYPE_INT64) {
        indicesInputTensor.dtype = "int64";
    } else {
        HCCL_ERROR("datatype[%s] is not int32 or int64", GetDataTypeEnumStr(dataType).c_str());
        return HCCL_E_PARA;
    }

    indicesInputTensor.shape = { static_cast<s64>(outputCount) };
    indicesInputTensor.oriShape = { static_cast<s64>(outputCount) };
    indicesInputTensor.format = "ND";
    indicesInputTensor.oriFormat = "ND";

    TeOpTensorArg indicesTilingArg;
    indicesTilingArg.tensor.push_back(indicesInputTensor);
    indicesTilingArg.argType = TensorArgType::TA_SINGLE;

    TeOpTensor outputTensor;
    outputTensor.dtype = "float32";
    outputTensor.shape = { static_cast<s64>(outputCount), valueDim };
    outputTensor.oriShape = { static_cast<s64>(outputCount), valueDim };
    outputTensor.format = "ND";
    outputTensor.oriFormat = "ND";

    TeOpTensorArg outputTilingArg;
    outputTilingArg.tensor.push_back(outputTensor);
    outputTilingArg.argType = TensorArgType::TA_SINGLE;

    TeOpParas paras;
    paras.inputs.push_back(inputTilingArg);
    paras.inputs.push_back(indicesTilingArg);
    paras.outputs.push_back(outputTilingArg);
    paras.attrs.push_back("0");

    GatherCompileInfos compileInfo;
    compileInfo.is_gather_v2 = opTilingInfo["compileInfo"]["is_gather_v2"];
    compileInfo.is_tik = opTilingInfo["compileInfo"]["is_tik"];
    compileInfo.core_num = opTilingInfo["compileInfo"]["vars"]["core_num"];
    compileInfo.impl_mode = opTilingInfo["compileInfo"]["vars"]["impl_mode"];
    compileInfo.indices_d_size = opTilingInfo["compileInfo"]["vars"]["indices_dsize"];
    compileInfo.is_preprocessed = opTilingInfo["compileInfo"]["vars"]["is_preprocessed"];
    compileInfo.l1_size = opTilingInfo["compileInfo"]["vars"]["l1_size"];
    compileInfo.params_d_size = opTilingInfo["compileInfo"]["vars"]["params_dsize"];
    compileInfo.ub_size = opTilingInfo["compileInfo"]["vars"]["ub_size"];

    string opName;
    CHK_RET(GetOpProperty(opDescInfo, "binFileName", opName)); // 获取json描述文件中数据
    bool isTiling = GatherTIKTiling(opName, paras, compileInfo, runInfo);
    CHK_PRT_RET(!isTiling, HCCL_ERROR("[Get][TilingRunInfo]gatherv2 tiling failed"), HCCL_E_CCE);

    return HCCL_SUCCESS;
}

// rt2.0
HcclResult TbeGatherV2::GetTilingRunInfo(nlohmann::json &opDescInfo,
    u64 inputCount, u64 outputCount, u32 valueDim, HcclDataType dataType, OpRunInfo &runInfo)
{
    ge::OpDescPtr op_desc = nullptr;
    EXECEPTION_CATCH(op_desc = make_shared<ge::OpDesc>("tmp_op", "gatherv2"), return HCCL_E_PTR);
    std::vector<int64_t> dimsIn = { static_cast<s64>(inputCount), valueDim };
    std::vector<int64_t> dimsInd = { static_cast<s64>(outputCount)};
    std::vector<int64_t> dimsOut = { static_cast<s64>(outputCount), valueDim};
    ge::GeTensorDesc tensor_descIn(ge::GeShape(dimsIn), ge::FORMAT_ND, ge::DataType::DT_FLOAT);
    ge::GeTensorDesc tensor_descInd(ge::GeShape(dimsInd), ge::FORMAT_ND, ge::DataType::DT_INT32);
    ge::GeTensorDesc tensor_descOut(ge::GeShape(dimsOut), ge::FORMAT_ND, ge::DataType::DT_FLOAT);

    op_desc->AddInputDesc("x", tensor_descIn);
    op_desc->AddInputDesc("indices", tensor_descInd);
    op_desc->AddOutputDesc("y", tensor_descOut);
    ge::ComputeGraphPtr graph = nullptr;
    EXECEPTION_CATCH(graph = make_shared<ge::ComputeGraph>("tmp"), return HCCL_E_PTR);
    ge::NodePtr node = graph->AddNode(op_desc);
    CHK_RET(AddDataNodeForOriginNode(graph, node, INPUT_SIZE));
    ge::Operator op = ge::OpDescUtils::CreateOperatorFromNode(node);

    // 解析
    std::string compileInfoJson = opDescInfo.dump();
    HCCL_DEBUG("compileInfoJson[%s]", compileInfoJson.c_str());
    auto tilingParseContext =
        gert::TilingParseContextBuilder()
            .CompileJson(compileInfoJson.c_str())
            .PlatformInfo(reinterpret_cast<void *>(const_cast<fe::PlatFormInfos *>
                (&g_gatherPlatFormInfos)))
            .CompileInfoCreatorFunc(CreateCompileInfo<optilingGather::GatherV2CompileInfo>)
            .CompileInfoDeleterFunc(DeleteCompileInfo<optilingGather::GatherV2CompileInfo>)
            .Build(op);

    auto parseContext = reinterpret_cast<gert::TilingParseContext *>(tilingParseContext.GetKernelContext());
    CHK_PTR_NULL(parseContext);

    ge::graphStatus gRet = optilingGather::TilingPrepareForGatherV2(parseContext);
    CHK_PRT_RET(gRet != ge::GRAPH_SUCCESS, HCCL_ERROR("[GetTilingRunInfo] Auto TilingParser failed"), HCCL_E_PARA);
    auto compiledInfo = parseContext->GetCompiledInfo<optilingGather::GatherV2CompileInfo>();

    // 构造tiling入参，进行tiling操作
    uint32_t maxSize;
    if (GetOpProperty(opDescInfo, "opParaSize", maxSize) != HCCL_SUCCESS) {
        HCCL_INFO("No max tiling size in opdesc.");
        maxSize = kMAX_TILING_DATA_SIZE;
    }
    const auto alignedMaxSize = RoundUp(static_cast<uint64_t>(maxSize), sizeof(uintptr_t));
    const auto tilingData = gert::TilingData::CreateCap(alignedMaxSize);
    const auto workspaceSize = gert::ContinuousVector::Create<size_t>(kMAX_WORKSPACE_COUNT);
    auto context_builder = gert::TilingContextBuilder();
    const gert::KernelContextHolder tilingContext =
        context_builder
            .CompileInfo(reinterpret_cast<void *>(compiledInfo))
            .PlatformInfo(reinterpret_cast<void *>(const_cast<fe::PlatFormInfos *>(&g_gatherPlatFormInfos)))
            .TilingData(tilingData.get())
            .Workspace(reinterpret_cast<gert::ContinuousVector *>(workspaceSize.get()))
            .Build(op);
    gert::TilingContext *context = reinterpret_cast<gert::TilingContext *>(tilingContext.context_);
    CHK_PTR_NULL(context);

    auto rawTilingData = context->GetRawTilingData();
    if (rawTilingData != nullptr) {
        rawTilingData->SetDataSize(0);
    }

    gRet = optilingGather::GatherTiling(context);
    CHK_PRT_RET(gRet != ge::GRAPH_SUCCESS, HCCL_ERROR("[GetTilingRunInfo] Auto Tiling failed"), HCCL_E_PARA);
    CHK_RET(ConvertToOpRunInfo(context, runInfo));
    return HCCL_SUCCESS;
}
#endif

HcclResult TbeGatherV2::Run(const void *embVal, const void *refMat, u64 inputCount,
    u64 outputCount, u32 valueDim, const HcclDataType dataType, void *stream, const void *dst)
{
    CHK_PRT_RET(!isInit_, HCCL_ERROR("[Run][Reduce]TbeGatherV2 had not been initialized"),
        HCCL_E_UNAVAIL);
    CHK_PTR_NULL(embVal);
    CHK_PTR_NULL(refMat);
    CHK_PTR_NULL(dst);
    CHK_PTR_NULL(stream);
    CHK_PRT_RET(inputCount == 0, HCCL_WARNING("input parameter 'inputCount' is 0"), HCCL_SUCCESS);
    CHK_PRT_RET(outputCount == 0, HCCL_WARNING("input parameter 'outputCount' is 0"), HCCL_SUCCESS);
    CHK_PRT_RET(dataType >= HCCL_DATA_TYPE_RESERVED,
        HCCL_ERROR("[Run][Reduce]unsupported data type[%s]", GetDataTypeEnumStr(dataType).c_str()), HCCL_E_PARA);

    indicesType_ = dataType;

    HCCL_DEBUG("inputCount[%llu], outputCount[%llu], valueDim[%u]", inputCount, outputCount, valueDim);

    // 获取op关联的描述信息和tiling信息
    nlohmann::json opDescInfo;
    nlohmann::json opTilingInfo;
    CHK_RET(GetOpInfo(opDescInfo, opTilingInfo));
    HCCL_INFO("opDescInfo[%s]", opDescInfo.dump().c_str());

    // 获取tiling data
    OpRunInfo runInfo;
    runInfo.blockDim = 0;
    runInfo.workspaces = { 0 };
    runInfo.clearAtomic = 0;
#ifndef HCCD
    if (SupportRt2()) {
        CHK_RET(GetTilingRunInfo(opDescInfo,
            inputCount, outputCount, valueDim, dataType, runInfo));
    } else {
        CHK_RET(GetTilingRunInfo(opDescInfo, opTilingInfo,
            inputCount, outputCount, valueDim, dataType, runInfo));
        // 申请overflow内存
        CHK_RET(hrtMalloc(&overFlowPtr_, MEMSET_OVERFLOW_SZIE));
    }
#else
        HCCL_ERROR("[TbeGatherV2]Does not support this interface.");
        return HCCL_E_NOT_SUPPORT;
#endif

    // 将tilingdata拷贝device侧
    unique_lock<mutex> tilingDataLock(tilingDataMutex_);
    void *tilingDataDevPtr = nullptr;
    CHK_RET(GetTilingDataDevMem(runInfo, &tilingDataDevPtr, MEMSET_GATHER_TILING_DATA_MAX_SZIE));
    tilingDataLock.unlock();

    // 构造桩函数
    string opName;
    CHK_RET(GetOpProperty(opDescInfo, "binFileName", opName));

    s8 *stubFunc = nullptr;
    string kernelName;
    CHK_RET(GetOpProperty(opDescInfo, "kernelName", kernelName));

    unique_lock<mutex> stubFuncLock(stubFuncMutex_);
    CHK_RET(GetstubFunc(opName, stubFunc, kernelName));
    stubFuncLock.unlock();

    // 执行OP
    CHK_RET(
        ExecuteKernelLaunch(embVal, refMat, dst, stream, axis_,
            reinterpret_cast<void *>(stubFunc), runInfo, tilingDataDevPtr, overFlowPtr_));

    return HCCL_SUCCESS;
}

HcclResult TbeGatherV2::GetChipName(std::string &chipVersion)
{
    s8 chipVer[CHIP_VERSION_MAX_LEN] = {0};
    CHK_RET(hrtGetSocVer(chipVer, CHIP_VERSION_MAX_LEN));
    char *chipName = reinterpret_cast<char*>(chipVer);
    std::string chipNameStr(chipName);
    auto iter = GATHERV2_SOC_STRING_MAP.find(chipNameStr);
    CHK_PRT_RET(iter == GATHERV2_SOC_STRING_MAP.end(),
        HCCL_ERROR("[Get][OpInfoIndex]chipVer[%s] is not supported",
        chipVer), HCCL_E_NOT_SUPPORT);
    chipVersion = iter->second;
    return HCCL_SUCCESS;
}

HcclResult TbeGatherV2::ExecuteKernelLaunch(const void *embVal, const void *refMat, const void *dst,
    void *stream, const void *axis, const void *stubFunc, OpRunInfo &runInfo, void *tilingDataDevPtr,
    const void *overFlowPtr)
{
    vector<const void *> deviceAddrs;
    deviceAddrs.push_back(embVal);
    deviceAddrs.push_back(refMat);
    deviceAddrs.push_back(axis);
    deviceAddrs.push_back(dst);
    deviceAddrs.push_back(tilingDataDevPtr);
    if (!SupportRt2()) {
        deviceAddrs.push_back(overFlowPtr);
    }

    HcclResult ret = hrtKernelLaunch(stubFunc, runInfo.blockDim, deviceAddrs.data(),
        deviceAddrs.size() * sizeof(void *), nullptr, stream);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[Execute][KernelLaunch]execute gatherv2 kernel launch failed"), HCCL_E_RUNTIME);

    return HCCL_SUCCESS;
}

HcclResult TbeGatherV2::GetopInfoIndex(string &opInfoIndex)
{
    if (indicesType_ == HCCL_DATA_TYPE_INT32) {
        switch (deviceType_) {
            case DevType::DEV_TYPE_910:
                opInfoIndex += "te_gatherv2_187b0f20dcf51283c20135753f6879d71370acd742114d5cb4847c86e660cd3c_1";
                break;
            case DevType::DEV_TYPE_310P3:
                opInfoIndex += "te_gatherv2_ec8c036da6e525f834dabb7a436b50387bbc5929f5201b28f0db7bd63fdb3a9b_1";
                break;
            case DevType::DEV_TYPE_910B:
            case DevType::DEV_TYPE_910_93:
                opInfoIndex += "gatherv2_Ascend910B_float32";
                break;
            default:
                HCCL_ERROR("deviceType_[%d] is not support", deviceType_);
                return HCCL_E_NOT_SUPPORT;
        }
    } else if (indicesType_ == HCCL_DATA_TYPE_INT64) {
        switch (deviceType_) {
            case DevType::DEV_TYPE_910:
                opInfoIndex += "te_gatherv2_08c676dd0564eaf646f0054fc0fb4adc559d7dca0155d698d2b788ce8a3a6a39_1";
                break;
            case DevType::DEV_TYPE_310P3:
                opInfoIndex += "te_gatherv2_ca87d1f41c037f24ece7993bc5c6f9afc00be851bcbaef7c03daf96701a30810_1";
                break;
            default:
                HCCL_ERROR("deviceType_[%d] is not support", deviceType_);
                return HCCL_E_NOT_SUPPORT;
        }
    } else {
        HCCL_ERROR("indicesType_[%d] is not int32 or int64", indicesType_);
        return HCCL_E_PARA;
    }

    return HCCL_SUCCESS;
}

HcclResult TbeGatherV2::GetCoreNum()
{
#ifndef HCCD
    s32 deviceLogicId;
    s64 coreNum;
    HcclRtDeviceModuleType hcclModuleType = HcclRtDeviceModuleType::HCCL_RT_DEVICE_MOUDLE_RESERVED;
    CHK_RET(hrtGetDevice(&deviceLogicId));

    if (deviceType_ == DevType::DEV_TYPE_910) {
        hcclModuleType = HcclRtDeviceModuleType::HCCL_RT_MODULE_TYPE_AICORE;
    } else if (deviceType_ == DevType::DEV_TYPE_910B || deviceType_ == DevType::DEV_TYPE_910_93) {
        hcclModuleType = HcclRtDeviceModuleType::HCCL_RT_MODULE_TYPE_VECTOR_CORE;
    } else {
        HCCL_ERROR("devType[%d] does not need tbe.", deviceType_);
        return HCCL_E_INTERNAL;
    }

    CHK_RET(hrtGetDeviceInfo(deviceLogicId, hcclModuleType, HcclRtDeviceInfoType::HCCL_INFO_TYPE_CORE_NUM, coreNum));
    HCCL_DEBUG("coreNum[%lld]", coreNum);
    g_gatherPlatFormInfos.SetCoreNum(static_cast<uint32_t>(coreNum));
#endif
    return HCCL_SUCCESS;
}

HcclResult TbeGatherV2::GetOpInfo(nlohmann::json &opDescInfo, nlohmann::json &opTilingInfo)
{
    string opInfoIndex;
    CHK_RET(GetopInfoIndex(opInfoIndex));
    if (SupportRt2()) {
        string chipNameStr;
        CHK_RET(GetChipName(chipNameStr));
        opInfoIndex += chipNameStr;
    }
    if (opInfoMap_.find(opInfoIndex) == opInfoMap_.end()) {
        HCCL_ERROR("[Get][OpInfo]can't find gather info,opInfoIndex[%s]", opInfoIndex.c_str());
        return HCCL_E_UNAVAIL;
    }
    opDescInfo = opInfoMap_[opInfoIndex];

    // rt1.0有tiling info
    if (!SupportRt2()) {
        opInfoIndex += "_tiling_info";
        if (opInfoMap_.find(opInfoIndex) == opInfoMap_.end()) {
            HCCL_ERROR("[Get][OpInfo]can't find gather info,opInfoTilingIndex[%s]", opInfoIndex.c_str());
            return HCCL_E_UNAVAIL;
        }
        opTilingInfo = opInfoMap_[opInfoIndex];
    }

    return HCCL_SUCCESS;
}

bool TbeGatherV2::SupportRt2()
{
    return deviceType_ == DevType::DEV_TYPE_910B || deviceType_ == DevType::DEV_TYPE_910_93;
}

void TbeGatherV2::InitOpInfoMap910A()
{
    opInfoMap_["te_gatherv2_187b0f20dcf51283c20135753f6879d71370acd742114d5cb4847c86e660cd3c_1"] =
        te_gatherv2_187b0f20dcf51283c20135753f6879d71370acd742114d5cb4847c86e660cd3c_1;
    opInfoMap_["te_gatherv2_187b0f20dcf51283c20135753f6879d71370acd742114d5cb4847c86e660cd3c_1_tiling_info"] =
        te_gatherv2_187b0f20dcf51283c20135753f6879d71370acd742114d5cb4847c86e660cd3c_1_tiling_info;
    opInfoMap_["te_gatherv2_08c676dd0564eaf646f0054fc0fb4adc559d7dca0155d698d2b788ce8a3a6a39_1"] =
        te_gatherv2_08c676dd0564eaf646f0054fc0fb4adc559d7dca0155d698d2b788ce8a3a6a39_1;
    opInfoMap_["te_gatherv2_08c676dd0564eaf646f0054fc0fb4adc559d7dca0155d698d2b788ce8a3a6a39_1_tiling_info"] =
        te_gatherv2_08c676dd0564eaf646f0054fc0fb4adc559d7dca0155d698d2b788ce8a3a6a39_1_tiling_info;
}

void TbeGatherV2::InitOpInfoMap310P3()
{
    opInfoMap_["te_gatherv2_ec8c036da6e525f834dabb7a436b50387bbc5929f5201b28f0db7bd63fdb3a9b_1"] =
        te_gatherv2_ec8c036da6e525f834dabb7a436b50387bbc5929f5201b28f0db7bd63fdb3a9b_1;
    opInfoMap_["te_gatherv2_ec8c036da6e525f834dabb7a436b50387bbc5929f5201b28f0db7bd63fdb3a9b_1_tiling_info"] =
        te_gatherv2_ec8c036da6e525f834dabb7a436b50387bbc5929f5201b28f0db7bd63fdb3a9b_1_tiling_info;
    opInfoMap_["te_gatherv2_ca87d1f41c037f24ece7993bc5c6f9afc00be851bcbaef7c03daf96701a30810_1"] =
        te_gatherv2_ca87d1f41c037f24ece7993bc5c6f9afc00be851bcbaef7c03daf96701a30810_1;
    opInfoMap_["te_gatherv2_ca87d1f41c037f24ece7993bc5c6f9afc00be851bcbaef7c03daf96701a30810_1_tiling_info"] =
        te_gatherv2_ca87d1f41c037f24ece7993bc5c6f9afc00be851bcbaef7c03daf96701a30810_1_tiling_info;
}

void TbeGatherV2::InitOpInfoMap910B()
{
    opInfoMap_["gatherv2_Ascend910B_float32_B1B2"] = gatherv2_Ascend910B_float32_B1B2;
    opInfoMap_["gatherv2_Ascend910B_float32_B3B4"] = gatherv2_Ascend910B_float32_B3B4;
}

} // namespace TbeReduce

#ifdef __cplusplus
extern "C" {
#endif

// 无业务验证使用
static TbeReduce::TbeGatherV2 g_gatherV2;
void Gatherv2Init()
{
    g_gatherV2.Init();
}

void Gatherv2Run(const void *embVal, const void *refMat, u64 inputCount,
    u64 outputCount, u32 valueDim, const HcclDataType dataType, void *stream, const void *dst)
{
    g_gatherV2.Run(embVal, refMat, inputCount, outputCount, valueDim, dataType, stream, dst);
}
#ifdef __cplusplus
}  // extern "C"
#endif
